import os
import sys
import torch
import numpy as np
import matplotlib.pyplot as plt
import cv2
import argparse

# Add current directory to Python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from config_loader import ConfigLoader
from model import SegmentationLightningModel

class Predictor:
    """Predictor class for image segmentation"""
    
    def __init__(self, config):
        self.config = config
        self.device = torch.device(config.DEVICE)
        
        # Create necessary directories
        config.create_directories()
        
        # Set random seed
        torch.manual_seed(42)
        np.random.seed(42)
    
    def load_model(self):
        """Load trained model"""
        print("Loading trained model...")
        
        # Find the best checkpoint
        checkpoint_files = [f for f in os.listdir(self.config.CHECKPOINT_DIR) 
                          if f.endswith('.ckpt')]
        
        if not checkpoint_files:
            print("Error: No checkpoint files found")
            return None
            
        # Use the first checkpoint found (prefer 'best' if available)
        best_checkpoints = [f for f in checkpoint_files if 'best' in f]
        if best_checkpoints:
            checkpoint_path = os.path.join(self.config.CHECKPOINT_DIR, best_checkpoints[0])
        else:
            checkpoint_path = os.path.join(self.config.CHECKPOINT_DIR, checkpoint_files[0])
        
        try:
            # Load model using PyTorch Lightning
            model = SegmentationLightningModel.load_from_checkpoint(
                checkpoint_path, config=self.config
            )
            print(f"Model loaded successfully: {checkpoint_path}")
            
            # Set model to evaluation mode
            model.eval()
            return model
        except Exception as e:
            print(f"Error loading model: {e}")
            return None

    def preprocess_image(self, image_path):
        """Preprocess input image"""
        # Read image
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Cannot load image: {image_path}")
        
        # Convert color space
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # Resize image
        image = cv2.resize(image, (self.config.IMAGE_SIZE[1], self.config.IMAGE_SIZE[0]))
        
        # Normalize
        image = image.astype(np.float32) / 255.0
        
        # Convert to tensor
        image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0)
        
        return image

    def predict_single_image(self, model, image_path, output_dir="output"):
        """Predict single image"""
        # Preprocess image
        image_tensor = self.preprocess_image(image_path)
        
        # Move tensor to the same device as model
        image_tensor = image_tensor.to(self.device)
        
        # Prediction
        with torch.no_grad():
            output = model(image_tensor)
            pred_mask = torch.argmax(output, dim=1).squeeze().cpu().numpy()
        
        # Create output directory
        os.makedirs(output_dir, exist_ok=True)
        
        # Save prediction results
        base_name = os.path.splitext(os.path.basename(image_path))[0]
        
        # Save mask
        mask_path = os.path.join(output_dir, f"{base_name}_mask.png")
        plt.imsave(mask_path, pred_mask, cmap='viridis')
        
        # Save visualization results
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
        
        # Original image
        original_image = cv2.imread(image_path)
        original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
        ax1.imshow(original_image)
        ax1.set_title("Original Image")
        ax1.axis('off')
        
        # Predicted mask
        ax2.imshow(pred_mask, cmap='viridis')
        ax2.set_title("Predicted Mask")
        ax2.axis('off')
        
        # Save visualization
        vis_path = os.path.join(output_dir, f"{base_name}_prediction.png")
        plt.savefig(vis_path, bbox_inches='tight', dpi=300)
        plt.close()
        
        print(f"Prediction saved: {vis_path}")
        return pred_mask

    def predict_batch_images(self, model, image_dir, output_dir="output"):
        """Predict batch of images"""
        # Get list of image files
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
        image_files = [f for f in os.listdir(image_dir) 
                      if os.path.splitext(f)[1].lower() in image_extensions]
        
        if not image_files:
            print(f"No images found in {image_dir}")
            return
        
        # Create output directory
        os.makedirs(output_dir, exist_ok=True)
        
        print(f"Predicting {len(image_files)} images...")
        
        for i, image_file in enumerate(image_files):
            image_path = os.path.join(image_dir, image_file)
            
            try:
                self.predict_single_image(model, image_path, output_dir)
                print(f"Processed {i+1}/{len(image_files)}: {image_file}")
            except Exception as e:
                print(f"Error processing {image_file}: {e}")

    def predict(self, image_path=None, image_dir=None, output_dir="output"):
        """Execute prediction"""
        print("=" * 60)
        print(f"Starting {self.config.MODEL_NAME} model prediction")
        print(f"Dataset: {self.config.DATASET_CONFIG.DATASET_NAME}")
        print("=" * 60)
        
        # Validate dataset paths
        try:
            self.config.DATASET_CONFIG.validate_paths()
            print(f"Dataset paths validated successfully")
        except FileNotFoundError as e:
            print(f"Warning: {e}")
            print(f"Dataset validation failed, but continuing with prediction...")
        
        # Load model
        model = self.load_model()
        if model is None:
            return
        
        # Prediction
        if image_path:
            self.predict_single_image(model, image_path, output_dir)
        elif image_dir:
            self.predict_batch_images(model, image_dir, output_dir)
        
        print("\nPrediction completed!")
        print("=" * 60)

def main():
    """Main function"""
    parser = argparse.ArgumentParser(description="Image segmentation prediction")
    parser.add_argument("--config", type=str, required=True, help="Configuration file path")
    parser.add_argument("--image", type=str, help="Single image path")
    parser.add_argument("--image_dir", type=str, help="Image directory path")
    parser.add_argument("--output", type=str, default="output", help="Output directory")
    
    args = parser.parse_args()
    
    if not args.image and not args.image_dir:
        print("Please specify --image or --image_dir parameter")
        return
    
    # Load configuration
    config = ConfigLoader.load_config(args.config)
    
    # Create predictor
    predictor = Predictor(config)
    
    # Execute prediction
    predictor.predict(args.image, args.image_dir, args.output)

if __name__ == "__main__":
    main()